package cn.johnyu.util;

import com.alibaba.dashscope.utils.JsonUtils;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class AliEmbeddingClient {
    protected static final Logger logger = LoggerFactory.getLogger(AliEmbeddingClient.class);
    public static List<List<Float>> embedding(List<String> texts) throws Exception {

        List<List<Float>> result = new ArrayList<>();
        String apiKey = System.getenv("DASHSCOPE_API_KEY");
        logger.info("获取到的API_KEY: "+apiKey);
        String baseUrl = "https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings";
        HttpClient client = HttpClient.newHttpClient();

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "text-embedding-v3");
        requestBody.put("input", texts);
        requestBody.put("dimensions", 1024);
        requestBody.put("encoding_format", "float");

        String requestBodyString = JsonUtils.toJson(requestBody);
        logger.info("请求体: "+requestBodyString);
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(baseUrl))
                .header("Content-Type", "application/json")
                .header("Authorization", "Bearer " + apiKey)
                .POST(HttpRequest.BodyPublishers.ofString(requestBodyString))
                .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (response.statusCode() == 200) {
            logger.info("请求成功，响应数据: " + response.body().substring(0,50)+" ......");
            //使用Gson完成解析，提取出数据：List<List<Float>>
            logger.info("提取数据，并转换成为向量列表...");
            result=extractEmbeddingFromResponse(response.body());
            logger.info("提取数据成功，向量列表大小: "+result.size());

        } else {
            logger.info("请求失败，状态码: "+response.statusCode());
        }
        return result;
    }
    //使用Gson完成解析，提取出数据：List<List<Float>>
    private static List<List<Float>>  extractEmbeddingFromResponse(String res) {
        List<List<Float>> result = new ArrayList<>();
        JsonObject object = JsonUtils.parse(res);
        JsonArray array = object.getAsJsonArray("data");
        for(JsonElement element:array){
            JsonArray vectorArr = element.getAsJsonObject().getAsJsonArray("embedding");
            List<Float> vector = new ArrayList<>();
            for (JsonElement jsonElement : vectorArr) {
                vector.add(jsonElement.getAsFloat());
            }
            result.add(vector);
        }
        return result;
    }

    public static void main(String[] args) throws Exception{
        List<String> texts = new ArrayList<>();
        texts.add("中华人民共和国刑法（第232条）: 故意杀人的，处死刑、无期徒刑或者十年以上有期徒刑；情节较轻的，处三年以上十年以下有期徒刑。");
        texts.add("中华人民共和国民法典（第123条）: 侵权责任法是调整侵权行为所产生的法律关系的法律。");
        texts.add("中华人民共和国公司法（第45条）: 股东会是公司的最高权力机构，股东会的决议对公司和全体股东具有约束力。");

        List<List<Float>> vectors = embedding(texts);


    }
}
